from Crypto.Cipher import AES
from Crypto.Util.Padding import pad, unpad
import base64
import os

# 密钥必须是16、24或32字节长
SECRET_KEY = os.environ.get('ENCRYPTION_KEY', 'my_secure_encryption_key_12345')
# 确保密钥长度正确
SECRET_KEY = SECRET_KEY.ljust(32)[:32].encode('utf-8')
IV = os.environ.get('ENCRYPTION_IV', 'initialization_vec').ljust(16)[:16].encode('utf-8')


def encrypt(data):
    """加密数据"""
    cipher = AES.new(SECRET_KEY, AES.MODE_CBC, IV)
    padded_data = pad(data.encode('utf-8'), AES.block_size)
    encrypted_data = cipher.encrypt(padded_data)
    return base64.b64encode(encrypted_data).decode('utf-8')


def decrypt(encrypted_data):
    """解密数据"""
    cipher = AES.new(SECRET_KEY, AES.MODE_CBC, IV)
    decoded_data = base64.b64decode(encrypted_data)
    decrypted_data = unpad(cipher.decrypt(decoded_data), AES.block_size)
    return decrypted_data.decode('utf-8')